package com.attrisoft.sanger.wrapper;

import java.io.PrintStream;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.SortedSet;
import java.util.TreeSet;

public class Matrix {
    
    private Map<String,Map<String,Float>> matrix = new HashMap<String,Map<String, Float>>();
    
    public void put(String columnKey, String rowKey, int value) {
        put(columnKey, rowKey, (float) value);
    }
    
    public void put(String columnKey, String rowKey, float value) {
        Map<String,Float> column = matrix.get(columnKey);
        if (column == null) {
            column = new HashMap<String,Float>();
            matrix.put(columnKey, column);
        }
        column.put(rowKey, value);
    }
    
    public float get(String columnKey, String rowKey) {
        Map<String,Float> column = matrix.get(columnKey);
        return (column == null) ? (Float) null : column.get(rowKey);
    }
    
    public Matrix getWeighted() {
        List<SortedSet<String>> keySets = getKeySets();
        SortedSet<String> columnKeys = keySets.get(0);
        SortedSet<String> rowKeys = keySets.get(1);
        Matrix weighted = new Matrix();
        for (String columnKey : columnKeys) {
            Map<String,Float> column = matrix.get(columnKey);
            if (column == null) {
                column = new HashMap<String,Float>();
                matrix.put(columnKey, column);
            }
            for (String rowKey : rowKeys) {
                float value = (float) 0;
                if (columnKey.equals(rowKey)) 
                    value = (float) 100;
                else if (column.get(rowKey) != null) {
                    value = (float) column.get(rowKey);
                }
                weighted.put(columnKey, rowKey, value);
            }
        }
        return weighted;
    }
    
    public Matrix getMarkov() {
        Matrix markov = new Matrix();
        for (String columnKey : matrix.keySet()) {
            Map<String,Float> column = matrix.get(columnKey);
            double sum = 0.0;
            for (Float value : column.values()) {
                sum += value;
            }
            for (String rowKey : column.keySet()) {
                float value = (float)(column.get(rowKey)/sum);
                markov.put(columnKey, rowKey, value);
            }
        }
        return markov;
    }
    
    public void toAbc(PrintStream out) {
        for (String columnKey : matrix.keySet()) {
            Map<String,Float> column = matrix.get(columnKey);
            for (String rowKey : column.keySet()) {
                out.print(columnKey +" " +rowKey +" ");
                out.println(String.valueOf(column.get(rowKey) == null ? 0 : column.get(rowKey)));
            }
        }
    }
    
    private List<SortedSet<String>> getKeySets() {
        SortedSet<String> columnKeys = new TreeSet<String>();
        SortedSet<String> rowKeys = new TreeSet<String>();
        for (String columnKey : matrix.keySet()) 
            columnKeys.add(columnKey);
        for (String columnKey : matrix.keySet()) {
            Map<String,Float> column = matrix.get(columnKey);
            for (String rowKey : column.keySet()) {
                rowKeys.add(rowKey);
            }
        }
        List<SortedSet<String>> keySets = new ArrayList<SortedSet<String>>();
        keySets.add(columnKeys);
        keySets.add(rowKeys);
        return keySets;
    }
    
    public void toCsv(PrintStream out) {
        List<SortedSet<String>> keySets = getKeySets();
        SortedSet<String> columnKeys = keySets.get(0);
        SortedSet<String> rowKeys = keySets.get(1);
        
        for (String columnKey : columnKeys) 
            out.print("," +columnKey);
        out.println();
        
        for (String rowKey : rowKeys) {
            out.print(rowKey);
            for (String columnKey : columnKeys) {
                if (matrix.get(columnKey).get(rowKey) == null) 
                    out.print(",0");
                else
                    out.print("," +matrix.get(columnKey).get(rowKey));
            }
            out.println();
        }
    }
    
    public void toRowIndex(PrintStream out) {
        List<SortedSet<String>> keySets = getKeySets();
        SortedSet<String> rowKeys = keySets.get(1);
        for (String rowKey : rowKeys) 
            out.println(rowKey);
    }
}
